package com.xjgz.aio.socket.enhance;

import java.io.IOException;
import java.net.SocketAddress;
import java.net.SocketOption;
import java.nio.ByteBuffer;
import java.nio.channels.AlreadyConnectedException;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.CompletionHandler;
import java.nio.channels.ConnectionPendingException;
import java.nio.channels.SelectionKey;
import java.nio.channels.ShutdownChannelGroupException;
import java.nio.channels.SocketChannel;
import java.util.Set;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

class EnhanceAsynchronousSocketChannel extends AsynchronousSocketChannel {

    /**
     * 底层的Socket通道，用于实际的网络IO操作
     * 该通道是非阻塞式的，支持异步读写操作
     */
    private final SocketChannel channel;

    /**
     * 处理读事件的工作线程，负责异步读取操作
     */
    private final EnhanceAsynchronousChannelGroup.Worker readWorker;

    /**
     * 读缓冲区
     */
    private ByteBuffer readBuffer;

    /**
     * 写缓冲区
     */
    private ByteBuffer writeBuffer;

    /**
     * 读操作的回调处理器，用于处理异步读取完成后的回调逻辑
     */
    private CompletionHandler<Number, Object> readCompletionHandler;

    /**
     * 写操作的回调处理器，用于处理异步写入完成后的回调逻辑
     */
    private CompletionHandler<Number, Object> writeCompletionHandler;

    /**
     * 读操作的附件对象，用于传递用户自定义数据
     */
    private Object readAttachMent;

    /**
     * 写操作的附件对象，用于传递用户自定义数据
     */
    private Object writeAttachment;

    /**
     * 读操作的SelectionKey对象，管理通道的读事件注册
     */
    private SelectionKey readSelectionKey;

    /**
     * 低内存模式
     */
    private final boolean lowMemory;

    /**
     * 写操作中断标记
     */
    private boolean writeInterrupted;

    /**
     * 读操作的调用次数，用于限制连续读取的次数，避免某个连接持续占用读取线程
     */
    private byte readInvoker = EnhanceAsynchronousChannelGroup.MAX_INVOKER;


    public EnhanceAsynchronousSocketChannel(
        EnhanceAsynchronousChannelGroup group,
        SocketChannel channel,
        boolean lowMemory
    ) {
        super(group.provider());
        this.channel = channel;
        this.readWorker = group.getReadWorker();
        this.lowMemory = lowMemory;
    }

    protected EnhanceAsynchronousChannelGroup group() {
        return readWorker.group();
    }

    @Override
    public void close() throws IOException {
        IOException exception = null;
        try {
            if (channel.isOpen()) {
                channel.close();
            }
        } catch (IOException e) {
            exception = e;
        }

        if (readCompletionHandler != null) {
            doRead(true);
        }
        if (readSelectionKey != null) {
            readSelectionKey.cancel();
            readSelectionKey = null;
        }
        SelectionKey key = channel.keyFor(group().writeWorker.selector);
        if (key != null) {
            key.cancel();
        }
        key = channel.keyFor(group().commnoWorker.selector);
        if (key != null) {
            key.cancel();
        }
        if (exception != null) {
            throw exception;
        }
    }

    @Override
    public boolean isOpen() {
        return channel.isOpen();
    }

    @Override
    public <T> T getOption(SocketOption<T> name) throws IOException {
        return channel.getOption(name);
    }

    @Override
    public Set<SocketOption<?>> supportedOptions() {
        return channel.supportedOptions();
    }

    @Override
    public AsynchronousSocketChannel bind(SocketAddress local) throws IOException {
        channel.bind(local);
        return this;
    }

    @Override
    public <T> AsynchronousSocketChannel setOption(SocketOption<T> name, T value) throws IOException {
        channel.setOption(name, value);
        return this;
    }

    @Override
    public AsynchronousSocketChannel shutdownInput() throws IOException {
        channel.shutdownInput();
        return this;
    }

    @Override
    public AsynchronousSocketChannel shutdownOutput() throws IOException {
        channel.shutdownOutput();
        return this;
    }

    @Override
    public SocketAddress getRemoteAddress() throws IOException {
        return channel.getRemoteAddress();
    }

    /**
     * 异步连接远程地址
     */
    @Override
    public <A> void connect(SocketAddress remote, A attachment, CompletionHandler<Void, ? super A> handler) {
        if (group().isTerminated()) {
            throw new ShutdownChannelGroupException();
        }
        if (channel.isConnected()) {
            throw new AlreadyConnectedException();
        }
        if (channel.isConnectionPending()) {
            throw new ConnectionPendingException();
        }
        doConnect(remote, attachment, handler);
    }

    @SuppressWarnings("rawtypes")
    private <A> void doConnect(SocketAddress remote, A attachment, CompletionHandler<Void, ? super A> handler) {
        try {
            if (handler instanceof FutureCompletionHandler && ((FutureCompletionHandler) handler).isDone()) {
                return;
            }
            boolean connected = channel.isConnectionPending();
            if (connected || channel.connect(remote)) {
                connected = channel.finishConnect();
            }
            channel.configureBlocking(false);
            if (connected) {
                handler.completed(null, attachment);
            } else {
                group().commnoWorker.addRegister(selector -> {
                    try {
                        channel.register(
                            selector,
                            SelectionKey.OP_CONNECT,
                            (Runnable) () -> doConnect(remote, attachment, handler)
                        );
                    } catch (ClosedChannelException e) {
                        handler.failed(e, attachment);
                    }
                });
            }
        } catch (IOException e) {
            handler.failed(e, attachment);
        }
    }

    @Override
    public Future<Void> connect(SocketAddress remote) {
        FutureCompletionHandler<Void, Void> connecFuture = new FutureCompletionHandler<>();
        connect(remote, null, connecFuture);
        return connecFuture;
    }

    @Override
    public <A> void read(
        ByteBuffer dst,
        long timeout,
        TimeUnit unit,
        A attachment,
        CompletionHandler<Integer, ? super A> handler
    ) {
        read0(dst, attachment, handler);
    }

    @SuppressWarnings("unchecked")
    private <V extends Number, A> void read0(
        ByteBuffer readBuffer,
        A attachment,
        CompletionHandler<V, ? super A> handler
    ) {
        this.readBuffer = readBuffer;
        this.readAttachMent = attachment;
        this.readCompletionHandler = (CompletionHandler<Number, Object>) handler;
        doRead(handler instanceof FutureCompletionHandler);
    }

    @Override
    public Future<Integer> read(ByteBuffer readBuffer) {
        FutureCompletionHandler<Integer, Object> readFuture = new FutureCompletionHandler<>();
        read(readBuffer, 0, TimeUnit.MILLISECONDS, null, readFuture);
        return readFuture;
    }

    @Override
    public <A> void read(
        ByteBuffer[] dsts,
        int offset,
        int length,
        long timeout,
        TimeUnit unit,
        A attachment,
        CompletionHandler<Long, ? super A> handler
    ) {
        throw new UnsupportedOperationException();
    }

    @Override
    public <A> void write(
        ByteBuffer src,
        long timeout,
        TimeUnit unit,
        A attachment,
        CompletionHandler<Integer, ? super A> handler
    ) {
        write0(src, attachment, handler);
    }

    @SuppressWarnings("unchecked")
    private <V extends Number, A> void write0(
        ByteBuffer writeBuffer,
        A attachment,
        CompletionHandler<V, ? super A> handler
    ) {
        this.writeBuffer = writeBuffer;
        this.writeAttachment = attachment;
        this.writeCompletionHandler = (CompletionHandler<Number, Object>) handler;
        while (doWrite());
    }

    @Override
    public Future<Integer> write(ByteBuffer src) {
        throw new UnsupportedOperationException();
    }

    @Override
    public <A> void write(
        ByteBuffer[] srcs,
        int offset,
        int length,
        long timeout,
        TimeUnit unit,
        A attachment,
        CompletionHandler<Long, ? super A> handler
    ) {
        throw new UnsupportedOperationException();
    }

    @Override
    public SocketAddress getLocalAddress() throws IOException {
        return channel.getLocalAddress();
    }

    /**
     * 异步读取操作
     */
    @SuppressWarnings("rawtypes")
    public final void doRead(boolean direct) {
        try {
            if (readCompletionHandler == null) {
                return;
            }

            if (readCompletionHandler instanceof FutureCompletionHandler && ((FutureCompletionHandler) readCompletionHandler)
                .isDone()
            ) {
                EnhanceAsynchronousChannelGroup.removeOps(readSelectionKey, SelectionKey.OP_READ);
                resetRead();
                return;
            }

            if (lowMemory && direct && readBuffer == null) {
                CompletionHandler<Number, Object> completionHandler = readCompletionHandler;
                Object attach = readAttachMent;
                resetRead();
                completionHandler.completed(EnhanceAsynchronousChannelProvider.READABLE_SIGNAL, attach);
                return;
            }

            boolean directRead = direct || readInvoker++ < EnhanceAsynchronousChannelGroup.MAX_INVOKER;

            int readSize = 0;
            boolean hasRemain = true;
            if (directRead) {
                readSize = channel.read(readBuffer);
                hasRemain = readBuffer.hasRemaining();
                if (hasRemain) {
                    readInvoker = EnhanceAsynchronousChannelGroup.MAX_INVOKER;
                }
            }

            if (readSize == 0 && readCompletionHandler instanceof FutureCompletionHandler) {
                EnhanceAsynchronousChannelGroup.removeOps(readSelectionKey, SelectionKey.OP_READ);
                group().commnoWorker.addRegister(selector -> {
                    try {
                        channel.register(selector, SelectionKey.OP_READ, EnhanceAsynchronousSocketChannel.this);
                    } catch (ClosedChannelException e) {
                        doRead(true);
                    }
                });
                return;
            }

            if (lowMemory && readSize == 0 && readBuffer.position() == 0) {
                readBuffer = null;
                readCompletionHandler.completed(EnhanceAsynchronousChannelProvider.READ_MONITOR_SIGNAL, readAttachMent);
            }

            if (readSize != 0 || !hasRemain) {
                CompletionHandler<Number, Object> completionHandler = readCompletionHandler;
                Object attach = readAttachMent;
                resetRead();
                completionHandler.completed(readSize, attach);

                if (readCompletionHandler == null && readSelectionKey != null) {
                    EnhanceAsynchronousChannelGroup.removeOps(readSelectionKey, SelectionKey.OP_READ);
                }
            } else if (readSelectionKey == null) {
                readWorker.addRegister(selector -> {
                    try {
                        if (channel.isOpen()) {
                            readSelectionKey =
                                channel.register(selector, SelectionKey.OP_READ, EnhanceAsynchronousSocketChannel.this);
                        }
                    } catch (Exception e) {
                        if (readCompletionHandler != null) {
                            readCompletionHandler.failed(e, readAttachMent);
                        }
                    }
                });
            } else {
                EnhanceAsynchronousChannelGroup.interestOps(readWorker, readSelectionKey, SelectionKey.OP_READ);
            }
        } catch (Exception e) {
            if (readCompletionHandler == null) {
                try {
                    close();
                } catch (Throwable ignore) {
                }
            } else {
                CompletionHandler<Number, Object> completionHandler = readCompletionHandler;
                Object attach = readAttachMent;
                resetRead();
                completionHandler.failed(e, attach);
            }
        } finally {
            readInvoker = 0;
        }
    }

    private void resetRead() {
        readCompletionHandler = null;
        readAttachMent = null;
        readBuffer = null;
    }

    public final boolean doWrite() {
        if (writeInterrupted) {
            writeInterrupted = false;
            return false;
        }
        try {
            int writeSize = channel.write(writeBuffer);

            if (writeSize != 0 || !writeBuffer.hasRemaining()) {
                CompletionHandler<Number, Object> handler = writeCompletionHandler;
                Object attachment = writeAttachment;
                resetWrite();
                writeInterrupted = true;
                handler.completed(writeSize, attachment);
                if (!writeInterrupted) {
                    return true;
                }
                writeInterrupted = false;
            } else {
                SelectionKey commSelectionKey = channel.keyFor(group().writeWorker.selector);
                if (commSelectionKey == null) {
                    group().writeWorker.addRegister(selector -> {
                        try {
                            if (channel.isOpen()) {
                                channel
                                    .register(selector, SelectionKey.OP_WRITE, EnhanceAsynchronousSocketChannel.this);
                            }
                        } catch (ClosedChannelException e) {
                            if (writeCompletionHandler != null) {
                                writeCompletionHandler.failed(e, writeAttachment);
                            }
                        }
                    });
                } else {
                    EnhanceAsynchronousChannelGroup
                        .interestOps(group().writeWorker, commSelectionKey, SelectionKey.OP_WRITE);
                }
            }
        } catch (Throwable e) {
            if (writeCompletionHandler == null) {
                e.printStackTrace();
                try {
                    close();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            } else {
                writeCompletionHandler.failed(e, writeAttachment);
            }
        }
        return false;
    }

    private void resetWrite() {
        writeCompletionHandler = null;
        writeAttachment = null;
        writeBuffer = null;
    }
}
